#include <algorithm>

#include "tinyxml.h"

#include "SpriteModule.h"

#define SPT_PTR_TABLE		(0x2f153f+4)
#define SPRITE_PALS_ADDR	0x30200

// a good place to write sprite gfx data, this is what the game uses
#define SPRITE_GFX_LOC_START	0x110200
#define SPRITE_GFX_LOC_END	0x1551FF

#define BUF_SIZE		6000

using namespace std;

Sprite::Sprite(uint w, uint h) {
  width = w;
  height = h;
}

Sprite::Sprite(uint newAddr, uint newWidth, uint newHeight, Rom &rom) {
  addr = newAddr;
  width = newWidth;
  height = newHeight;
  uint i, j;
  uint offset = 0;
  image = new uchar[(width*8)*(height*8)];
  for (i=0; i<width*height*8*8; i++)
    image[i] = 0;
  for (i=0; i<height; i++) {
    for (j=0; j<width; j++) {
      read4BppImage(&((rom.getBuffer())[addr+offset]), image, j*8, i*8, width*8, 0);
      offset += 32;
    }
  }
}

uint Sprite::writeToRom(uint address, Rom& rom) {
  uint i, j;
  uint offset = 0;
  addr = address;
  for (i=0; i<height; i++) {
    for (j=0; j<width; j++) {
      write4BppImage(image, &((rom.getBuffer())[addr+offset]), j*8, i*8, width*8, 0);
      offset += 32;
    }
  }
  return offset;
}

uint Sprite::getRomSize() {
  return 32*width*height;
}

Sprite::~Sprite() {
  delete[] image;
}

uint Sprite::getAddr() {
  return addr;
}

uint Sprite::getWidth() {
  return width;
}

uint Sprite::getHeight() {
  return height;
}

void Sprite::getAsString(char* out) {
  if (width*height*8*8 >= BUF_SIZE) {
    cerr << "Error: sprite too big for buffer" << endl;
    out[0] = '\0';
    return;
  }
  uint i;
  for (i=0; i<(width*8)*(height*8); i++) {
    itoa(image[i], out++, 16);
  }
  out[0] = '\0';
}

void Sprite::setAsString(const char* str) {
  int i, len = width*height*8*8;
  image = new uchar[len];
  for (i=0; i<len; i++)
    image[i] = (uchar) chtoi(str[i]);
}

SpritePointer::SpritePointer() {
  sprite = 0;
  flipped = false;
}

SpritePointer::SpritePointer(uint spr, bool flip) {
  sprite = spr;
  flipped = flip;
}

void SpritePointer::setSprite(uint spr) {
  sprite = spr;
}

uint SpritePointer::getSprite() {
  return sprite;
}

void SpritePointer::setFlipped(bool flip) {
  flipped = flip;
}

bool SpritePointer::isFlipped() {
  return flipped;
}

void SpritePointer::load(uint addr, uint width, uint height, vector<Sprite*> &sprites, Rom &rom) {
  flipped = ((addr & 1) == 1);
  addr &= 0xfffffe;

  // See if the sprite was already loaded
  int i = 0;
  vector<Sprite*>::iterator it;
  for (it = sprites.begin(); it != sprites.end(); it++) {
    if ((*it)->getAddr() == addr) {
      sprite = i;
      return;
    }
    i++;
  }

  // Load the new sprite
  Sprite* sp = new Sprite(addr, width, height, rom);
  sprites.push_back(sp);
  sprite = i;
}

uint SpritePointer::getPointer(vector<Sprite*> &sprites) {
  return (toSnesPointer(sprites.at(sprite)->getAddr()) | (flipped ? 1 : 0)) & 0xffff;
}

SpriteSet::SpriteSet() {
  pointers = 0;
}

SpriteSet::~SpriteSet() {
  if (pointers)
    delete[] pointers;
}

uint SpriteSet::getRomLength() {
  return 9 + 2*num_sprites;
}

void SpriteSet::writeToRom(uint addr, vector<Sprite*> &sprites, Rom &rom) {
  int i;
  uint bank;
  rom.write(height, addr++);
  rom.write(width<<4, addr++);
  rom.write(unknown[0], addr++);
  rom.write(palette<<1, addr++);
  for(i=1; i<5; i++)
    rom.write(unknown[i], addr++);
  bank = toSnesPointer(sprites.at(pointers[0].getSprite())->getAddr()) >> 16;
  rom.write(bank, addr++);
  for (i=0; i<num_sprites; i++) {
    rom.writeMulti(pointers[i].getPointer(sprites), addr, 2);
    addr += 2;
  }
}

void SpritePalette::readFromRom(uint addr, Rom &rom) {
  int i;
  uint bgrBlock;
  for (i=0; i<16; i++) {
    bgrBlock = rom.readMulti(addr+i*2,2) & 0x7fff;
    colors[i].r = (bgrBlock & 0x001f) * 8;
    colors[i].g = ((bgrBlock & 0x03e0) >> 5) * 8;
    colors[i].b = (bgrBlock >> 10) * 8;
  }
}

void SpritePalette::writeToRom(uint addr, Rom &rom) {
  int i;
  uint bgrBlock;
  for (i=0; i<16; i++) {
    bgrBlock = (colors[i].r >> 3) & 0x1f;
    bgrBlock |= ((colors[i].g >> 3) & 0x1f) << 5;
    bgrBlock |= ((colors[i].b >> 3) & 0x1f) << 10;
    bgrBlock &= 0x7fff;

    rom.writeMulti(bgrBlock, addr + 2*i, 2);
  }
}

void SpritePalette::getAsString(char* out) {
  int i;
  for (i=0; i<16; i++) {
    itoa(colors[i].r >> 3, out++, 32);
    itoa(colors[i].g >> 3, out++, 32);
    itoa(colors[i].b >> 3, out++, 32);
  }
}

void SpritePalette::setAsString(const char* str) {
  int i;
  for (i=0; i<16; i++) {
    colors[i].r = chtoi(*(str++)) << 3;
    colors[i].g = chtoi(*(str++)) << 3;
    colors[i].b = chtoi(*(str++)) << 3;
  }
}

BankGroup::BankGroup() {
  size = 0;
}

uint BankGroup::getSize() {
  return size;
}

bool BankGroup::sharesSprites(SpriteSet& ss, vector<Sprite*> &sprites) {
  int i;
  for (i=0; i<ss.num_sprites; i++) {
    if (hasSprite(sprites.at(ss.pointers[i].getSprite())))
      return true;
  }
  return false;
}

bool BankGroup::hasSprite(Sprite* sprite) {
  return (find(bank.begin(), bank.end(), sprite) != bank.end());
}

void BankGroup::addSpriteSet(SpriteSet& ss, vector<Sprite*> &sprites) {
  int i;
  Sprite* sp;
  for (i=0; i<ss.num_sprites; i++) {
    sp = sprites.at(ss.pointers[i].getSprite());
    if (!hasSprite(sp)) {
      bank.push_back(sp);
      size += sp->getRomSize();
    }
  }
}

vector<Sprite*>* BankGroup::getVector() {
  return &bank;
}

bool isSpriteBigger(Sprite* sp1, Sprite* sp2) {
  return (sp1->getWidth() * sp1->getHeight()) > (sp2->getWidth() * sp2->getHeight());
}

void BankGroup::sortDescending() {
  sort(bank.begin(), bank.end(), isSpriteBigger);
}

Bank::Bank() {
  addr = 0;
  free = 0;
}

Bank::Bank(uint ad, uint capacity) {
  addr = ad;
  free = capacity;
}

void Bank::init(uint address, uint size) {
  addr = address;
  free = size;
}

bool Bank::hasRoomFor(BankGroup* bg) {
  return (bg->getSize() <= free);
}

void Bank::add(BankGroup* bg) {
  vector<Sprite*>* sps = bg->getVector();
  vector<Sprite*>::iterator it;
  free -= bg->getSize();
  for (it = sps->begin(); it != sps->end(); it++)
    bank.push_back(*it);
}

void Bank::writeToRom(Rom &rom) {
//  cout << "Writing bank @ 0x" << hex << addr << dec << ", " << free << " bytes unused" << endl;
  uint writeAddr = addr;
  vector<Sprite*>::iterator it;
  for (it = bank.begin(); it != bank.end(); it++) {
    writeAddr += (*it)->writeToRom(writeAddr, rom);
  }
}

SpriteModule::~SpriteModule() {
  vector<Sprite*>::iterator it;
  for (it = sprites.begin(); it != sprites.end(); it++)
    delete (*it);
  sprites.clear();
}

const char* SpriteModule::getName() {
  return "Sprites";
}

const char* SpriteModule::getCredits() {
  return "Written by MrTenda";
}

void SpriteModule::readFromRom(Rom &rom) {
  int i, j;
  uint addr, nextAddr, bank;

  // Read palettes
  for (i=0; i < NUM_PALS; i++)
    palettes[i].readFromRom(SPRITE_PALS_ADDR + 32*i, rom);

  // Read sprites
  addr = toRegPointer(rom.readMulti(SPT_PTR_TABLE, 4));
  for (i=0; i < NUM_SPRITE_SETS; i++) {
    spriteSets[i].height = rom.read(addr++);
    spriteSets[i].width = rom.read(addr++) >> 4;
    spriteSets[i].unknown[0] = rom.read(addr++);
    spriteSets[i].palette = (rom.read(addr++)>>1) & 0x7;
    for (j=1; j < 5; j++) {
      spriteSets[i].unknown[j] = rom.read(addr++);
    }

    bank = toRegPointer(rom.read(addr++)<<16);
    if (i == NUM_SPRITE_SETS - 1) {
      // Note: can we be more intelligent about this??
      spriteSets[i].num_sprites = 8;
    } else {
      nextAddr = toRegPointer(rom.readMulti(SPT_PTR_TABLE + ((i+1)*4), 4));
      spriteSets[i].num_sprites = (nextAddr - addr) / 2;
    }
    spriteSets[i].pointers = new SpritePointer[spriteSets[i].num_sprites];

    for (j=0; j<spriteSets[i].num_sprites; j++) {
      spriteSets[i].pointers[j].load(
        bank + rom.readMulti(addr, 2),
        spriteSets[i].width, spriteSets[i].height,
        sprites, rom);
      addr += 2;
    }
  }
}

void SpriteModule::writeToRom(Rom &rom) { 
  int i;
  uint addr, len;
  bool tmp;

  // #1: Write palettes to ROM
//  cout << "Writing palettes..." << endl;
  for (i=0; i < NUM_PALS; i++)
    palettes[i].writeToRom(SPRITE_PALS_ADDR + 32*i, rom);

  // #2: Write sprite graphics to ROM
  // Note: this is complex because sprites in the same group need to be in the same rom bank
  // A "Bank Group" is a group of sprites required to be in the same bank
//  cout << "Creating bank groups..." << endl;
  vector<BankGroup*> bank_groups;
  BankGroup* bank_group;
  vector<BankGroup*>::iterator banks_it;
  bool found;
  for (i=0; i<NUM_SPRITE_SETS; i++) {
    // Check if any of the bank groups contain any pre-existing Sprites
    found = false;
    for (banks_it = bank_groups.begin(); banks_it != bank_groups.end(); banks_it++) {
      if ((*banks_it)->sharesSprites(spriteSets[i], sprites)) {
        (*banks_it)->addSpriteSet(spriteSets[i], sprites);
        found = true;
        break;
      }
      if (found)
        break;
    }
    // Have to create a new bank group
    if (!found) {
      bank_group = new BankGroup();
      bank_group->addSpriteSet(spriteSets[i], sprites);
      bank_groups.push_back(bank_group);
    }
  }

//  cout << "Created " << bank_groups.size() << " bank groups from " << NUM_SPRITE_SETS << " sprite sets." << endl;

//  cout << "Sorting bank groups..." << endl;
  // Bin packing problem... just use the First Fit Decreasing heuristic
  // We must pack the BankGroups into the Banks
  // First, sort the Bank Groups...
  for (banks_it = bank_groups.begin(); banks_it != bank_groups.end(); banks_it++)
    (*banks_it)->sortDescending(); 
  // Create the Bank objects
  Bank banks[5];
  i=0;
  for (addr=SPRITE_GFX_LOC_START; addr < SPRITE_GFX_LOC_END; addr += 0x10000) {
    if (addr + 0x10000 > SPRITE_GFX_LOC_END)
      banks[i++].init(addr, SPRITE_GFX_LOC_END - addr);
    else
      banks[i++].init(addr, 0x10000);
  }
  // Bin pack the bank groups into the banks, first fit
  int j=0;
  for (banks_it = bank_groups.begin(); banks_it != bank_groups.end(); banks_it++) {
    tmp = false;
    j++;
    for (i=0; i<5; i++) {
      if (banks[i].hasRoomFor(*banks_it)) {
        banks[i].add(*banks_it);
        delete (*banks_it);
        tmp = true;
        break;
      }
    }
    if (!tmp)
      cout << "Error: could not find room for bank group #" << j << " in any banks" << endl;
  }

  // Write banks to ROM
  for (i=0; i<5; i++) {
    banks[i].writeToRom(rom);
  }

  // #3: Now write the SPT entries to ROM
  // we have to write these contiguously so they can be read correctly later...
  len = 0;
  for (i=0; i<NUM_SPRITE_SETS; i++)
    len += spriteSets[i].getRomLength();
  // Find a good place to write the SPT
  addr = rom.getFreeRange(len);
//  cout << "Writing SPT to 0x" << hex << addr << dec << " (" << len << " bytes)" << endl;
  // Make SPT entry #0 point to entry #1
  rom.writeMulti(toSnesPointer(addr), SPT_PTR_TABLE - 4, 4);
  for (i=0; i<NUM_SPRITE_SETS; i++) {
    // Write SPT entry
    spriteSets[i].writeToRom(addr, sprites, rom);
    // Write pointer to SPT entry
    rom.writeMulti(toSnesPointer(addr), SPT_PTR_TABLE + (i*4), 4);
    addr += spriteSets[i].getRomLength();
  }
}

void SpriteModule::readFromDat(char* dirName, map<string, string> &fns, map<string, int> &labels) {
  char buf[100];
  getDatFname(fns, "sprites", "sprites.xml", dirName, buf);

  TiXmlDocument doc(buf);
  if (!doc.LoadFile()) {
    cerr << "Error: (" << getName() << ") Failed to load " << buf << endl;
    return;
  }
  TiXmlHandle hDoc(&doc);
  TiXmlHandle hRoot(0), hSection(0), hEntry(0), hField(0);
  TiXmlElement *entry, *subfield;

  uint i, j;

  hRoot = TiXmlHandle(hDoc.FirstChildElement().Element());

  // Read the palettes
  i=0;
  hSection = TiXmlHandle(hRoot.FirstChild("Palettes").Element());
  for (entry = hSection.FirstChild().Element(); entry; entry = entry->NextSiblingElement()) {
    i = atoi_prefix(entry->Attribute("id"));
    palettes[i].setAsString(entry->GetText());
  }

  // Read the sprite sets
  hSection = TiXmlHandle(hRoot.FirstChild("SpriteSets"));
  for (entry = hSection.FirstChild().Element(); entry; entry = entry->NextSiblingElement()) {
    hEntry = TiXmlHandle(entry);
    i = atoi_prefix(entry->Attribute("id"))-1;
    spriteSets[i].width = atoi_prefix(hEntry.FirstChild("Width").Element()->GetText());
    spriteSets[i].height = atoi_prefix(hEntry.FirstChild("Height").Element()->GetText());
    spriteSets[i].palette = atoi_prefix(hEntry.FirstChild("Palette").Element()->GetText());

    // Read the SpritePointers
    hField = TiXmlHandle(hEntry.FirstChild("Sprites").Element());
    j = 0;
    // First, count how many SpritePointers there are
    for (subfield = hField.FirstChild().Element(); subfield; subfield = subfield->NextSiblingElement()) {
      j++;
    }
    // Then create the array
    spriteSets[i].num_sprites = j;
    spriteSets[i].pointers = new SpritePointer[j];
    // Read in the array...
    j = 0;
    for (subfield = hField.FirstChild().Element(); subfield; subfield = subfield->NextSiblingElement()) {
      spriteSets[i].pointers[j].setSprite(atoi_prefix(subfield->GetText()));
      spriteSets[i].pointers[j].setFlipped(atoi_prefix(subfield->Attribute("flip")) != 0);
      j++;
    }

    // Read the unknowns
    hField = TiXmlHandle(hEntry.FirstChild("Unknowns").Element());
    j=0;
    for (subfield = hField.FirstChild().Element(); subfield; subfield = subfield->NextSiblingElement()) {
      spriteSets[i].unknown[j++] = atoi_prefix(subfield->GetText());
    }
  }

  // Read the sprite graphics
  hSection = TiXmlHandle(hRoot.FirstChild("Sprites"));
  Sprite* sp;
  for (entry = hSection.FirstChild().Element(); entry; entry = entry->NextSiblingElement()) {
    sp = new Sprite(
      atoi_prefix(entry->Attribute("width")),
      atoi_prefix(entry->Attribute("height")));
    sp->setAsString(entry->GetText());
    sprites.push_back(sp);
  }
}

void SpriteModule::writeToDat(char* dirName, map<string, string> &fns) {
  TiXmlDocument doc;
  TiXmlDeclaration* decl = new TiXmlDeclaration( "1.0", "", "" );
  doc.LinkEndChild(decl);

  TiXmlElement* root = new TiXmlElement("Sprites");
  doc.LinkEndChild(root);

  TiXmlElement *section, *entry, *subentry, *field;
  int i, j;
  char buf[BUF_SIZE];

  // Write the palettes
  section = new TiXmlElement("Palettes");
  root->LinkEndChild(section);
  for (i=0; i<NUM_PALS; i++) {
    entry = new TiXmlElement("Palette");
    entry->SetAttribute("id", i);
    palettes[i].getAsString(buf);
    entry->LinkEndChild( new TiXmlText( buf ));
    section->LinkEndChild(entry);
  }

  // Write the sprite sets
  section = new TiXmlElement("SpriteSets");
  root->LinkEndChild(section);
  for (i=0; i<NUM_SPRITE_SETS; i++) {
    entry = new TiXmlElement("SpriteSet");
    section->LinkEndChild(entry);
    entry->SetAttribute("id", i+1);
    field = new TiXmlElement("Width");
    entry->LinkEndChild(field);
    itoa(spriteSets[i].width, buf, 10);
    field->LinkEndChild(new TiXmlText(buf));
    field = new TiXmlElement("Height");
    entry->LinkEndChild(field);
    itoa(spriteSets[i].height, buf, 10);
    field->LinkEndChild(new TiXmlText(buf));
    field = new TiXmlElement("Palette");
    entry->LinkEndChild(field);
    itoa(spriteSets[i].palette, buf, 10);
    field->LinkEndChild(new TiXmlText(buf));

    field = new TiXmlElement("Sprites");
    entry->LinkEndChild(field);
    for (j=0; j<spriteSets[i].num_sprites; j++) {
      subentry = new TiXmlElement("Sprite");
      field->LinkEndChild(subentry);
      subentry->SetAttribute("flip", spriteSets[i].pointers[j].isFlipped());
      itoa(spriteSets[i].pointers[j].getSprite(), buf, 10);
      subentry->LinkEndChild(new TiXmlText(buf));
    }

    field = new TiXmlElement("Unknowns");
    entry->LinkEndChild(field);
    for (j=0; j<5; j++) {
      subentry = new TiXmlElement("Unknown");
      field->LinkEndChild(subentry);
      itoa_prefix(spriteSets[i].unknown[j], buf, 16);
      subentry->LinkEndChild(new TiXmlText(buf));
    }
  }

  // Write the sprites
  section = new TiXmlElement("Sprites");
  root->LinkEndChild(section);
  i=0;
  vector<Sprite*>::iterator it;
  for (it = sprites.begin(); it != sprites.end(); it++) {
    entry = new TiXmlElement("Sprite");
    entry->SetAttribute("id", i++);
    entry->SetAttribute("width", (*it)->getWidth());
    entry->SetAttribute("height", (*it)->getHeight());
    (*it)->getAsString(buf);
    entry->LinkEndChild(new TiXmlText(buf));
    section->LinkEndChild(entry);
  }

  getDatFname(fns, "sprites", "sprites.xml", dirName, buf);
  doc.SaveFile(buf);
}
